package net.csiit.platform.codeguard.classloader;

import net.csiit.platform.codeguard.tools.AES;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.JarURLConnection;
import java.net.URL;
import java.net.URLClassLoader;
import java.net.URLConnection;
import java.security.AccessController;
import java.security.PrivilegedExceptionAction;
import java.util.Set;
import java.util.jar.JarFile;

/**
 * @author adam
 * @date 2020/7/2
 */
public class DecryptClassLoader extends URLClassLoader {

    protected Set<String> decryptPackages;
    protected String decryptKey;

    public DecryptClassLoader(URL[] urls, ClassLoader parent, Set<String> decryptPackages, String decryptKey) {
        super(urls, parent);
        this.decryptPackages = decryptPackages;
        this.decryptKey = decryptKey;
    }

    public DecryptClassLoader(URL[] urls, Set<String> decryptPackages, String decryptKey) {
        super(urls);
        this.decryptPackages = decryptPackages;
        this.decryptKey = decryptKey;
    }

    public void setDecryptPackages(Set<String> decryptPackages) {
        this.decryptPackages = decryptPackages;
    }

    public void setDecryptKey(String decryptKey) {
        this.decryptKey = decryptKey;
    }

    @Override
    protected Class<?> findClass(String name) throws ClassNotFoundException {
        Class cls = findLoadedClass(name);
        if (cls != null) {
            return cls;
        }

        if (!CodeGuardConfiguration.matches(name)) {
            return super.findClass(name);
        }

        String path = name.replace('.', '/').concat(".class");

        InputStream inputStream = getParent() == null ? super.getResourceAsStream(path) : getParent().getResourceAsStream(path);
        if (inputStream == null) {
            throw new ClassNotFoundException(name);
        }

        try {
            try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream()) {
                byte[] buffer = new byte[4096];
                int readByte = -1;
                while ((readByte = inputStream.read(buffer)) != -1) {
                    byteArrayOutputStream.write(buffer, 0, readByte);
                }
                inputStream.close();

                byte[] bytes = byteArrayOutputStream.toByteArray();
                bytes = AES.decrypt(bytes, AES.pwdHandler(decryptKey));
                if (bytes == null) {
                    throw new ClassNotFoundException("Decrypt class failed !");
                }

                definePackageIfNecessary(name);
                return defineClass(name, bytes, 0, bytes.length);
            }
        } catch (IOException e) {
            throw new ClassNotFoundException("Cannot load resource for class [" + name + "]", e);
        }
    }

    /**
     * Define a package before a {@code findClass} call is made. This is necessary to
     * ensure that the appropriate manifest for nested JARs is associated with the
     * package.
     * @param className the class name being found
     */
    protected void definePackageIfNecessary(String className) {
        int lastDot = className.lastIndexOf('.');
        if (lastDot >= 0) {
            String packageName = className.substring(0, lastDot);
            if (getPackage(packageName) == null) {
                try {
                    definePackage(className, packageName);
                }
                catch (IllegalArgumentException ex) {
                    // Tolerate race condition due to being parallel capable
                    if (getPackage(packageName) == null) {
                        // This should never happen as the IllegalArgumentException
                        // indicates that the package has already been defined and,
                        // therefore, getPackage(name) should not have returned null.
                        throw new AssertionError(
                                "Package " + packageName + " has already been defined but it could not be found");
                    }
                }
            }
        }
    }

    private void definePackage(String className, String packageName) {
        try {
            AccessController.doPrivileged((PrivilegedExceptionAction<Object>) () -> {
                String packageEntryName = packageName.replace('.', '/') + "/";
                String classEntryName = className.replace('.', '/') + ".class";
                for (URL url : getURLs()) {
                    try {
                        URLConnection connection = url.openConnection();
                        if (connection instanceof JarURLConnection) {
                            JarFile jarFile = ((JarURLConnection) connection).getJarFile();
                            if (jarFile.getEntry(classEntryName) != null && jarFile.getEntry(packageEntryName) != null
                                    && jarFile.getManifest() != null) {
                                definePackage(packageName, jarFile.getManifest(), url);
                                return null;
                            }
                        }
                    }
                    catch (IOException ex) {
                        // Ignore
                    }
                }
                return null;
            }, AccessController.getContext());
        }
        catch (java.security.PrivilegedActionException ex) {
            // Ignore
        }
    }
}
